# coding=utf-8

import torch

# class UNIZARNet(torch.nn.Module):
#     def __init__(self, input_param_count: int):
#         super(UNIZARNet, self).__init__()
#         self.net = torch.nn.Sequential(
#             torch.nn.BatchNorm1d(input_param_count),
#             torch.nn.Sigmoid(),
#             torch.nn.Linear(input_param_count, 48),
#             torch.nn.BatchNorm1d(48),
#             torch.nn.ReLU(),
#             torch.nn.Linear(48, 64),
#             torch.nn.BatchNorm1d(64),
#             torch.nn.ReLU(),
#             torch.nn.Linear(64, 36),
#             torch.nn.BatchNorm1d(36),
#             torch.nn.ReLU(),
#             torch.nn.Linear(36, 24),
#             torch.nn.BatchNorm1d(24),
#             torch.nn.ReLU(),
#             torch.nn.Linear(24, 7),
#         )

#     def forward(self, input: torch.FloatTensor):
#         return self.net(input)

# def main():
#     eval_net = UNIZARNet(9)
#     model_save_path = "D:/Learn/solar_software/pytorch_script/unizar_net3-500.pth"
#     eval_net.load_state_dict(torch.load(model_save_path))
#     dummy_input = torch.randn(1024, 9)
#     onnx_program = torch.onnx.dynamo_export(eval_net, dummy_input)
#     onnx_program.save("unizar_net3-500.onnx")

import torch
import torch.nn as nn
import torch.nn.functional as F


class MyModel(nn.Module):

    def __init__(self):
        super(MyModel, self).__init__()
        self.conv1 = nn.Conv2d(1, 6, 5)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = torch.flatten(x, 1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x


# class UNIZARNet(torch.nn.Module):
#     def __init__(self, input_param_count: int):
#         super(UNIZARNet, self).__init__()
#         self.__input_param_count = input_param_count
#         # self.__running_mean1 = torch.zeros(self.__input_param_count, dtype=torch.float32)
#         # self.__running_var1 = torch.ones(self.__input_param_count, dtype=torch.float32)
#         # self.__running_mean2 = torch.zeros(48, dtype=torch.float32)
#         # self.__running_var2 = torch.ones(48, dtype=torch.float32)
#         # self.__running_mean3 = torch.zeros(48, dtype=torch.float32)
#         # self.__running_var3 = torch.ones(48, dtype=torch.float32)
#         self.__batch_norm1 = torch.nn.BatchNorm1d(48)
#         self.__batch_norm1.run
#         self.__linear1 = torch.nn.Linear(input_param_count, 48)
#         self.__linear2 = torch.nn.Linear(48, 64)
#         self.__linear3 = torch.nn.Linear(64, 36)
#         self.__linear4 = torch.nn.Linear(36, 24)
#         self.__linear5 = torch.nn.Linear(24, 7)

#     def forward(self, input: torch.Tensor):
#         input = F.batch_norm(input, None, None)
#         input = F.sigmoid(input)
#         # input = F.batch_norm(input, self.__running_mean2, self.__running_var2)
#         input = F.relu(F.batch_norm(self.__linear1(input), None, None))

#         return input


class UNIZARNet(torch.nn.Module):
    def __init__(self, input_param_count: int):
        super(UNIZARNet, self).__init__()
        self.net = torch.nn.Sequential(
            torch.nn.BatchNorm1d(input_param_count),
            torch.nn.Sigmoid(),
            torch.nn.Linear(input_param_count, 48),
            torch.nn.BatchNorm1d(48),
            torch.nn.ReLU(),
            torch.nn.Linear(48, 64),
            torch.nn.BatchNorm1d(64),
            torch.nn.ReLU(),
            torch.nn.Linear(64, 36),
            torch.nn.BatchNorm1d(36),
            torch.nn.ReLU(),
            torch.nn.Linear(36, 24),
            torch.nn.BatchNorm1d(24),
            torch.nn.ReLU(),
            torch.nn.Linear(24, 7),
        )

    def forward(self, input: torch.FloatTensor):
        return self.net(input)


def main():
    torch_model = UNIZARNet(9)
    torch_model.load_state_dict(
        torch.load("D:/Learn/solar_software/pytorch_script/unizar_net3-500.pth")
    )
    torch_input = torch.randn(2, 9)
    # onnx_program = torch.onnx.dynamo_export(torch_model, torch_input)
    onnx_program = torch.onnx.export(
        torch_model,
        torch_input,
        input_names=["input"],
        output_names=["output"],
        dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}},
        f="unizar_net3-500.onnx",
    )
    # onnx_program.save("my_image_classifier.onnx")


if __name__ == "__main__":
    main()
